#include "postprocess/kalman_filter.h"

bool KalmanFilter::s_use_adaptive_ = true;
bool KalmanFilter::use_ssld_ = true;
bool KalmanFilter::s_use_breakdown_ = true;
float KalmanFilter::s_association_score_maximum_ = 8.0;
Eigen::Matrix3f KalmanFilter::s_propagation_noise_ =
    10 * Eigen::Matrix3f::Identity();
float KalmanFilter::s_measurement_noise_ = 0.4;
float KalmanFilter::s_initial_velocity_noise_ = 5.0;
float KalmanFilter::s_initial_acceleration_noise_ = 0.6;
float KalmanFilter::s_predict_variance_per_sqrsec_ = 10.0;
float KalmanFilter::s_breakdown_threshold_maximum_ = 10;
uint32_t KalmanFilter::s_measurement_cached_history_size_minimum_ = 3;
uint32_t KalmanFilter::s_measurement_cached_history_size_maximum_ = 6;

KalmanFilter::KalmanFilter() {
  age_ = 0;
  measurement_cached_history_size_ = s_measurement_cached_history_size_minimum_;
  s_measurement_cached_history_size_maximum_ =
      s_measurement_cached_history_size_minimum_;
  velocity_covariance_ =
      s_initial_velocity_noise_ * Eigen::Matrix3f::Identity();
  // to do: add acceleration
  acceleration_covariance_ =
      s_initial_acceleration_noise_ * Eigen::Matrix3f::Identity();
  update_quality_ = 1.0;
  breakdown_threshold_ = s_breakdown_threshold_maximum_;
  belief_velocity_ = Eigen::Vector3f::Zero();
  belief_acceleration_gain_ = Eigen::Vector3f::Zero();
  belief_acceleration_ = Eigen::Vector3f::Zero();
  state_ = Eigen::Vector4f::Zero();
  measurement_covariance_ = Eigen::Matrix4f::Zero();
  state_covariance_ = Eigen::Matrix4f::Zero();
}

KalmanFilter::~KalmanFilter() {}

void KalmanFilter::Initialize(Eigen::Vector3f const& center_point,
                              Eigen::Vector3f const& velocity) {
  update_quality_ = 1.0;
  breakdown_threshold_ = s_breakdown_threshold_maximum_;
  belief_center_point_ = center_point.cast<float>();
  belief_velocity_ = velocity.cast<float>();
  belief_acceleration_ = Eigen::Vector3f::Zero();
  belief_acceleration_gain_ = Eigen::Vector3f::Zero();
  state_covariance_ = Eigen::Matrix4f::Identity();
  state_covariance_.block<2, 2>(0, 0) *= s_initial_velocity_noise_;
  state_covariance_.block<2, 2>(2, 2) *= s_initial_acceleration_noise_;
}

void KalmanFilter::GetAccelerationCovariance(
    Eigen::Matrix3f* acceleration_covariance) {
  (*acceleration_covariance) = acceleration_covariance_.cast<float>();
}

void KalmanFilter::GetAccelerationGain(Eigen::Vector3f* acceleration_gain) {
  (*acceleration_gain) = belief_acceleration_gain_.cast<float>();
}

void KalmanFilter::GetVelocityCovariance(Eigen::Matrix3f* velocity_covariance) {
  (*velocity_covariance) = velocity_covariance_.cast<float>();
}

void KalmanFilter::GetStateCovariance(Eigen::Matrix4f* state_covariance) {
  (*state_covariance) = state_covariance_.cast<float>();
}

void KalmanFilter::UpdateWithoutObject(double time_diff) {
  // only update belief anchor point
  belief_center_point_ += belief_velocity_ * time_diff;
  age_ += 1;
}

void KalmanFilter::UpdateWithObject(TrackedObjectPtr& new_object,
                                    const TrackedObjectPtr& old_object,
                                    double time_diff) {
  // 1. compute update quality if needed
  if (s_use_adaptive_) {
    compute_update_quality(new_object, old_object);
  } else {
    update_quality_ = 1.0;
  }

  // 2. compute motion measurement
  select_measured_velocity_ =
      compute_velocity(new_object, old_object, time_diff);
  select_measured_acceleration_ = new_object->acceleration;

  // 3. kalman filter update
  update_state(new_object, old_object, time_diff);

  // 4. cache measurement history  and update belief
  if (history_measured_velocity_.size() >= measurement_cached_history_size_) {
    history_measured_velocity_.pop_front();
    history_time_diff_.pop_front();
  }
  history_measured_velocity_.push_back(select_measured_velocity_);
  history_time_diff_.push_back(time_diff);

  // 5. Evaluate online covariance
  evaluate_online_covariance();
  age_ += 1;
}

void KalmanFilter::GetState(Eigen::Vector3f* anchor_point,
                            Eigen::Vector3f* velocity,
                            Eigen::Vector3f* acceleration) {
  (*anchor_point) = belief_center_point_.cast<float>();
  (*velocity) = belief_velocity_.cast<float>();
  (*acceleration) = belief_acceleration_.cast<float>();
}

void KalmanFilter::GetOnlineCovariance(Eigen::Matrix3f* online_covariance) {
  *online_covariance = online_velocity_covariance_.cast<float>();
}

void KalmanFilter::update_state(const TrackedObjectPtr& new_object,
                                const TrackedObjectPtr& old_object,
                                double time_diff) {
  // 1. rotate composition
  Eigen::Matrix4f transition = Eigen::Matrix4f::Identity();
  transition(0, 2) = transition(1, 3) = static_cast<float>(time_diff);

  // 2. prediction stage
  // tmp fix for 12b
  Eigen::Matrix4f predict_covariance = Eigen::Matrix4f::Identity();
  //  predict_covariance(0, 0) *= 0.08;
  //  predict_covariance(1, 1) *= 0.30;
  //  predict_covariance(2, 2) *= 0.02;
  //  predict_covariance(3, 3) *= 0.05;

  predict_covariance(0, 0) *= 8;
  predict_covariance(1, 1) *= 30;
  predict_covariance(2, 2) *= 2;
  predict_covariance(3, 3) *= 5;
  // tmp fix for 12b
  // LOG4CXX_DEBUG(logger_, "state_: " << state_);
  state_ = transition * state_;
  state_covariance_ = transition * state_covariance_ * transition.transpose() +
                      predict_covariance;

  // 3. measurement update stage
  Eigen::Matrix2f measurement_covariance = Eigen::Matrix2f::Identity();
  Eigen::Vector2f measurement;
  measurement << select_measured_velocity_.head<2>();
  // Eigen::Vector2f new_dir = new_object->direction.head<2>();
  // new_dir.normalize();
  // Eigen::Vector2f odir(new_dir(1), -new_dir(0));
  measurement_covariance(0, 0) *= 0.55;
  measurement_covariance(1, 1) *= 0.18;
  // tmp fix for 12b
  Eigen::Matrix<float, 2, 4> observation_transform;
  observation_transform.block<2, 2>(0, 0).setIdentity();
  observation_transform.block<2, 2>(0, 2).setZero();
  Eigen::Matrix<float, 4, 2> kalman_gain_matrix =
      static_cast<Eigen::Matrix<float, 4, 2, 0, 4, 2>>(
          state_covariance_ * observation_transform.transpose() *
          (observation_transform * state_covariance_ *
               observation_transform.transpose() +
           measurement_covariance)
              .inverse());
  Eigen::Vector4f state_gain = static_cast<Eigen::Matrix<float, 4, 1, 0, 4, 1>>(
      kalman_gain_matrix * (measurement - observation_transform * state_));
  // state gain adjustment and estimate posterior
  state_gain_adjustment(state_gain);
  // add state gain and update state covariance
  state_ = state_ + state_gain;
  state_covariance_ = (Eigen::Matrix4f::Identity() -
                       kalman_gain_matrix * observation_transform) *
                      state_covariance_;
  belief_center_point_ = new_object->center;
  belief_velocity_.head<2>() = state_.head<2>();
  belief_velocity_(2) = 0.0;
  belief_acceleration_.head<2>() = state_.tail<2>();
  belief_acceleration_(2) = 0.0;
}

void KalmanFilter::evaluate_online_covariance() {
  Eigen::Matrix3f online_covariance = Eigen::Matrix3f::Zero();
  std::int32_t evaluate_window =
      history_measured_velocity_.size() >
              s_measurement_cached_history_size_maximum_
          ? history_measured_velocity_.size()
          : s_measurement_cached_history_size_maximum_;
  for (std::int32_t i = 0; i < evaluate_window; ++i) {
    int history_index = history_measured_velocity_.size() - i - 1;
    Eigen::Vector3f velocity_resisual = Eigen::Vector3f(5, 5, 0);
    if (history_index >= 0) {
      velocity_resisual =
          history_measured_velocity_[history_index] - belief_velocity_;
    }
    online_covariance(0, 0) += velocity_resisual(0) * velocity_resisual(0);
    online_covariance(0, 1) += velocity_resisual(0) * velocity_resisual(1);
    online_covariance(1, 0) += velocity_resisual(1) * velocity_resisual(0);
    online_covariance(1, 1) += velocity_resisual(1) * velocity_resisual(1);
  }
  if (evaluate_window != 0) {
    online_velocity_covariance_ = online_covariance / evaluate_window;
  } else {
    online_covariance = Eigen::Matrix3f::Zero();
  }
}

void KalmanFilter::state_gain_adjustment(Eigen::Vector4f& gain) {
  if (s_use_adaptive_) {
    gain *= update_quality_;
  }
  // 2. breakdown the constrained the manx allowed change of state
  if (s_use_breakdown_) {
    // velocity breakdown threshold
    float velocity_breakdown_threshold = 0.3;
    float velocity_gain = gain.head<2>().norm();
    if (velocity_gain > velocity_breakdown_threshold) {
      gain.head<2>() *= velocity_breakdown_threshold / velocity_gain;
    }
    // acceleration breakdown threshold
    float acceleration_breakdown_threshold = 1.0;
    float acceleration_gain = gain.tail<2>().norm();
    if (acceleration_gain > acceleration_breakdown_threshold) {
      gain.tail<2>() *= acceleration_breakdown_threshold / acceleration_gain;
    }
  }
}

void KalmanFilter::compute_update_quality(const TrackedObjectPtr& new_object,
                                          const TrackedObjectPtr& old_object) {
  // Compute update quality for adaptive filtering
  // according to association score
  float update_quality_according_association_score =
      compute_update_quality_according_association_score(new_object);
  update_quality_ = update_quality_according_association_score;
}

float KalmanFilter::compute_update_quality_according_association_score(
    const TrackedObjectPtr& new_object) {
  float association_score = new_object->association_score;
  float update_quality = 1;
  if (s_association_score_maximum_ == 0) {
    return update_quality;
  }
  update_quality = 1 - association_score / s_association_score_maximum_;
  update_quality = update_quality > 1 ? 1 : update_quality;
  update_quality = update_quality < 0 ? 0 : update_quality;
  update_quality = update_quality * update_quality;
  return update_quality;
}

Eigen::VectorXf KalmanFilter::compute_velocity(
    TrackedObjectPtr const& new_object, TrackedObjectPtr const& old_object,
    double time_diff) noexcept {
  // center point velocity
  Eigen::Vector3f center_velocity =
      compute_measured_center_velocity(new_object, old_object, time_diff);
  // corner point velocity
  Eigen::Vector3f corner_velocity =
      compute_measured_corner_velocity(new_object, old_object, time_diff);
  std::vector<Eigen::Vector3f> measured_candidates;
  measured_candidates.push_back(center_velocity);
  measured_candidates.push_back(corner_velocity);

  // select measured velocity among candidates according motion consistency
  // todo @kening consider direction consistency as well
  // measure x, y directory seperately
  Eigen::Vector3f old_velocity = old_object->velocity;
  Eigen::Vector3f select_measured_velocity;
  int candidates_size = measured_candidates.size();
  float select_measured_x_gain_velocity = FLT_MAX;
  float select_measured_y_gain_velocity = FLT_MAX;
  int select_measured_x_ind{0};
  int select_measured_y_ind{0};
  if (candidates_size > 0) {
    // x dir
    for (int i = 0; i < candidates_size; ++i) {
      float measured_x_velocity_gain =
          measured_candidates[i].x() - old_velocity.x();
      if (fabs(measured_x_velocity_gain) <
          fabs(select_measured_x_gain_velocity)) {
        select_measured_x_gain_velocity = measured_x_velocity_gain;
        select_measured_x_ind = i;
      }
    }
    // y dir
    for (int i = 0; i < candidates_size; ++i) {
      float measured_y_velocity_gain =
          measured_candidates[i].y() - old_velocity.y();
      if (fabs(measured_y_velocity_gain) <
          fabs(select_measured_y_gain_velocity)) {
        select_measured_y_gain_velocity = measured_y_velocity_gain;
        select_measured_y_ind = i;
      }
    }
    select_measured_velocity << measured_candidates[select_measured_x_ind].x(),
        measured_candidates[select_measured_y_ind].y(), 0;
  } else {
    select_measured_velocity = old_velocity;
  }
  return select_measured_velocity;
}

Eigen::VectorXf KalmanFilter::compute_measured_center_velocity(
    TrackedObjectPtr const& new_object, TrackedObjectPtr const& old_object,
    double time_diff) noexcept {
  Eigen::Vector3f measured_center_point_velocity =
      (new_object->center - old_object->center) / static_cast<float>(time_diff);
  measured_center_point_velocity(2) = 0.0;
  return measured_center_point_velocity;
}

Eigen::VectorXf KalmanFilter::compute_measured_corner_velocity(
    TrackedObjectPtr const& new_object, TrackedObjectPtr const& old_object,
    double time_diff) noexcept {
  Eigen::Vector3f old_bbox_corner_list[4], new_bbox_corner_list[4];
  for (std::int16_t i = 0; i < 4; ++i) {
    old_bbox_corner_list[i].x() = old_object->polygon.Get(i).x;
    old_bbox_corner_list[i].x() = old_object->polygon.Get(i).y;
    old_bbox_corner_list[i].z() = 0.0f;
    new_bbox_corner_list[i].x() = new_object->polygon.Get(i).x;
    new_bbox_corner_list[i].y() = new_object->polygon.Get(i).y;
    new_bbox_corner_list[i].z() = 0.0f;
  }
  // match old bbox corner and new bbox corner by order
  float min_total_shift{FLT_MAX};
  int order[4] = {0, 1, 2, 3};
  for (int i{0}; i < 4; ++i) {
    float tmp_total_shift{0};
    for (int j{0}; j < 4; ++j) {
      float shift =
          (old_bbox_corner_list[j] - new_bbox_corner_list[j]).squaredNorm();
      tmp_total_shift += shift;
    }
    if (tmp_total_shift < min_total_shift) {
      min_total_shift = tmp_total_shift;
      // shift order i
      if (i > 0) {
        ReverseArray(order, 0, i - 1);
        ReverseArray(order, i, 3);
        ReverseArray(order, 0, 3);
        Eigen::Vector3f tmp_bbox_corner_list[4];
        for (std::uint8_t k{0}; k < 4; k++) {
          tmp_bbox_corner_list[k] = new_bbox_corner_list[order[k]];
        }
        for (std::uint8_t k{0}; k < 4; k++) {
          new_bbox_corner_list[k] = tmp_bbox_corner_list[k];
        }
      }
    }
  }
  // check cloeset corner point
  float closest_dis{FLT_MAX};
  Eigen::Vector3f old_closest_corner_point;
  std::uint8_t old_closest_corner_point_order{0};
  for (std::uint8_t i{0}; i < 4; ++i) {
    float tmp_dis = old_bbox_corner_list[i].squaredNorm();
    if (tmp_dis < closest_dis) {
      closest_dis = tmp_dis;
      old_closest_corner_point = old_bbox_corner_list[i];
      old_closest_corner_point_order = i;
    }
  }
  // check whether last frame and current frame has the same closest point
  closest_dis = FLT_MAX;
  Eigen::Vector3f new_closest_corner_point;
  std::uint8_t new_closest_corner_point_order{0};
  for (std::uint8_t i{0}; i < 4; ++i) {
    float tmp_dis = new_bbox_corner_list[i].squaredNorm();
    if (tmp_dis < closest_dis) {
      closest_dis = tmp_dis;
      new_closest_corner_point = new_bbox_corner_list[i];
      new_closest_corner_point_order = i;
    }
  }
  if (old_closest_corner_point_order == new_closest_corner_point_order) {
    // no closest point shifting
    Eigen::Vector3f bbox_corner_velocity =
        ((new_closest_corner_point - old_closest_corner_point) /
         static_cast<float>(time_diff))
            .cast<float>();
    bbox_corner_velocity(2) = 0.0;
  } else {
    // closest corner point has shiftted
    new_closest_corner_point =
        new_bbox_corner_list[old_closest_corner_point_order];
    Eigen::Vector3f bbox_corner_velocity =
        ((new_closest_corner_point - old_closest_corner_point) /
         static_cast<float>(time_diff))
            .cast<float>();
    bbox_corner_velocity(2) = 0.0;
  }
  Eigen::Vector3f min_bbox_corner_velocity = Eigen::Vector3f(100.0, 100.0, 0);
  float min_bbox_corner_velocity_gain_norm = min_bbox_corner_velocity.norm();
  for (size_t i = 0; i < 4; ++i) {
    Eigen::Vector3f bbox_corner_velocity =
        ((new_bbox_corner_list[i] - old_bbox_corner_list[i]) / time_diff)
            .cast<float>();
    bbox_corner_velocity(2) = 0.0;
    if (bbox_corner_velocity.norm() < min_bbox_corner_velocity_gain_norm) {
      min_bbox_corner_velocity = bbox_corner_velocity;
      min_bbox_corner_velocity_gain_norm = bbox_corner_velocity.norm();
    }
  }
  return min_bbox_corner_velocity;
}

void KalmanFilter::ReverseArray(int arr[], int begin, int end) noexcept {
  int temp;
  for (; begin < end; begin++, end--) {
    temp = arr[end];
    arr[end] = arr[begin];
    arr[begin] = temp;
  }
}
